import random


class Label:
    def __init__(self, file_id, chinese, pinyin):
        self.file_id = file_id
        self.chinese = chinese
        self.pinyin = pinyin


def read_labels(filename: str) -> list:
    with open(filename, encoding='utf-8') as f:
        lines = f.readlines()
    assert len(lines) > 0 and len(lines) % 2 == 0

    labels = []
    for idx, line in enumerate(lines):
        if idx % 2 == 0:
            file_id, chinese = line.strip().split('\t')
        else:
            pinyin = line.strip()
            labels.append(Label(file_id, chinese, pinyin))

    return labels


def write_labels(filename: str, labels: list):
    with open(filename, 'w', encoding='utf-8') as f:
        for label in labels:
            f.write('DUMMY/{}.wav|{}\n'.format(label.file_id, label.pinyin))


if __name__ == '__main__':
    dist = {
        'test': 400,
        'train': 9500,
        'val': 100,
    }

    labels = read_labels('label.txt')
    random.shuffle(labels)
    for ty, num in dist.items():
        curr_slice = labels[:num]
        labels = labels[num:]
        write_labels('{}.txt'.format(ty), curr_slice)
